import pandas as pd
from sklearn.ensemble import GradientBoostingRegressor
from sklearn.model_selection import GridSearchCV, train_test_split
from sklearn.metrics import mean_squared_error
from sklearn.ensemble import StackingRegressor
from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor


def hyper_parameter_selection(model,param_grid, x_train: pd.DataFrame, y_train: pd.DataFrame):

    # 初始化GridSearchCV
    grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=5, scoring='neg_mean_squared_error', verbose=1)

    # 执行网格搜索
    grid_search.fit(x_train, y_train)

    # 输出最佳参数
    print(f"Best parameters: {grid_search.best_params_}")


def gbdt(train: pd.DataFrame, target: pd.DataFrame, test_x: pd.DataFrame, test_y: pd.DataFrame):
    x_train, x_val, y_train, y_val = train_test_split(train, target, test_size=0.2, random_state=2020)
    # 划分训练集和测试集

    # 初始化GBDT回归模型
    gbdt_regressor = GradientBoostingRegressor(n_estimators=200, learning_rate=0.2, max_depth=5, random_state=42)

    # 训练模型
    gbdt_regressor.fit(x_train, y_train)

    # 预测验证集
    y_pred = gbdt_regressor.predict(x_val)

    # 计算均方误差
    mse = mean_squared_error(y_val, y_pred)
    print(f"Mean Squared Error: {mse}")

    # 如果需要，可以输出模型的参数
    print(f"Model parameters: {gbdt_regressor.get_params()}")

    y_pred = gbdt_regressor.predict(test_x)
    mse = mean_squared_error(test_y, y_pred)
    print(f"Mean Squared Error: {mse}")


def stacking(train: pd.DataFrame, target: pd.DataFrame, test_x: pd.DataFrame, test_y: pd.DataFrame):
    X_train, x_val, y_train, y_val = train_test_split(train, target, test_size=0.2, random_state=2020)
    # 定义基模型
    estimators = [
        ('lr', LinearRegression()),
        ('dt', DecisionTreeRegressor(max_depth=4)),
        ('rf', RandomForestRegressor(n_estimators=10, max_depth=4)),
        ('gb', GradientBoostingRegressor(n_estimators=10, max_depth=3))
    ]

    # 定义元模型
    final_estimator = LinearRegression()

    # 初始化StackingRegressor
    stacking_regressor = StackingRegressor(estimators=estimators, final_estimator=final_estimator, cv=5)

    # 训练模型
    stacking_regressor.fit(X_train, y_train)

    # 预测测试集
    y_pred = stacking_regressor.predict(x_val)

    # 计算均方误差
    mse = mean_squared_error(y_val, y_pred)
    print(f"Mean Squared Error: {mse}")

    y_pred = stacking_regressor.predict(test_x)
    mse = mean_squared_error(test_y, y_pred)
    print(f"Mean Squared Error: {mse}")


input = pd.read_csv(r'.//new_input.csv', index_col=0)
index = pd.read_csv(r'.//AEI.csv', index_col=0)
input.drop(['advertiser_id', 'promotion_id'], inplace=True, axis=1)
train, x_test, target, y_test = train_test_split(input, index, test_size=0.2, random_state=2020)
gbdt(train, target, x_test, y_test)
# # 定义GBDT回归模型
# gbdt_regressor = GradientBoostingRegressor(random_state=42)
# 定义要搜索的超参数网格
# param_grid = {
#     'n_estimators': [50, 100, 200],
#     'learning_rate': [0.01, 0.1, 0.2],
#     'max_depth': [3, 4, 5],
#
# }
# hyper_parameter_selection(gbdt_regressor, param_grid, train, target)
# stacking(train, target, x_test, y_test)
